#include "core/enhanced_chat_system_db.h"
#include "core/common.h"
#include <signal.h>
#include <iostream>
#include <thread>
#include <vector>
#include <algorithm>

class EnhancedChatServer {
private:
    int serverSocket;
    int port;
    std::vector<std::shared_ptr<ClientInfo>> clients;
    std::mutex clientsMutex;
    std::atomic<bool> running;
    
    // 数据库持久化的聊天系统
    std::unique_ptr<EnhancedChatSystemDB> chatSystem;

public:
    EnhancedChatServer(int p = DEFAULT_PORT) : port(p), running(false) {
        chatSystem = std::make_unique<EnhancedChatSystemDB>();
    }
    
    ~EnhancedChatServer() {
        stop();
    }
    
    bool start() {
        // 初始化数据库连接
        DatabaseConfig dbConfig;
        dbConfig.host = "127.0.0.1";
        dbConfig.username = "root";
        dbConfig.password = "";
        dbConfig.database = "enhanced_chat_db";
        
        if (!chatSystem->initialize(dbConfig)) {
            printLog("ERROR: Failed to initialize database connection");
            return false;
        }
        
        // 创建socket
        serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            printLog("ERROR: Failed to create socket");
            return false;
        }
        
        // 设置socket选项
        int opt = 1;
        if (setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
            printLog("ERROR: Failed to set socket options");
            return false;
        }
        
        // 绑定地址
        sockaddr_in serverAddr;
        serverAddr.sin_family = AF_INET;
        serverAddr.sin_addr.s_addr = INADDR_ANY;
        serverAddr.sin_port = htons(port);
        
        if (bind(serverSocket, (sockaddr*)&serverAddr, sizeof(serverAddr)) < 0) {
            printLog("ERROR: Failed to bind socket to port " + std::to_string(port));
            return false;
        }
        
        // 开始监听
        if (listen(serverSocket, MAX_CLIENTS) < 0) {    
            printLog("ERROR: Failed to listen on socket");
            return false;
        }
        
        running = true;
        printLog("Enhanced chat server with database started on port " + std::to_string(port));
        printLog("Database connection established");
        printLog("Waiting for connections...");
        
        return true;
    }
    
    void run() {
        while (running) {
            sockaddr_in clientAddr;
            socklen_t clientLen = sizeof(clientAddr);
            
            int clientSocket = accept(serverSocket, (sockaddr*)&clientAddr, &clientLen);
            if (clientSocket < 0) {
                if (running) {
                    printLog("ERROR: Failed to accept client connection");
                }
                continue;
            }
            
            std::string clientIP = inet_ntoa(clientAddr.sin_addr);
            printLog("New connection from " + clientIP);
            
            // 为每个客户端创建处理线程
            std::thread clientThread(&EnhancedChatServer::handleClient, this, clientSocket, clientIP);
            clientThread.detach();
        }
    }
    
    void stop() {
        running = false;
        if (serverSocket >= 0) {
            close(serverSocket);
        }
        
        // 关闭所有客户端连接
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (auto client : clients) {
            close(client->socket);
        }
        clients.clear();
        
        // 清理聊天系统和数据库连接
        if (chatSystem) {
            chatSystem->cleanup();
        }
        
        printLog("Enhanced server stopped");
    }
    
private:
    void handleClient(int clientSocket, const std::string& clientIP) {
        char buffer[BUFFER_SIZE];
        std::string username;
        bool authenticated = false;
        
        while (running) {
            memset(buffer, 0, BUFFER_SIZE);
            int bytesReceived = recv(clientSocket, buffer, BUFFER_SIZE - 1, 0);
            
            if (bytesReceived <= 0) {
                break;
            }
            
            std::string data(buffer);
            
            // 尝试解析增强消息格式
            EnhancedMessage enhancedMsg = EnhancedMessage::deserialize(data);
            
            if (!authenticated && enhancedMsg.type == MSG_LOGIN) {
                username = enhancedMsg.fromUser;
                if (authenticateUser(username, clientSocket, clientIP)) {
                    authenticated = true;
                    addClient(clientSocket, username, clientIP);
                    
                    // 发送登录成功消息
                    EnhancedMessage loginSuccess(MSG_LOGIN, "Server", "Login successful");
                    sendToClient(clientSocket, loginSuccess);
                    
                    // 广播用户加入消息
                    EnhancedMessage joinMsg(MSG_USER_STATUS, "Server", username + " joined the chat");
                    broadcastMessage(joinMsg, clientSocket);
                    
                    // 发送在线用户列表
                    sendUserList(clientSocket);
                    
                    printLog("User '" + username + "' logged in from " + clientIP);
                } else {
                    EnhancedMessage loginFailed(MSG_ERROR, "Server", "Username already taken or database error");
                    sendToClient(clientSocket, loginFailed);
                }
            } else if (authenticated) {
                // 处理各种消息类型
                handleEnhancedMessage(enhancedMsg, username, clientSocket);
            }
        }
        
        // 清理
        if (authenticated) {
            chatSystem->removeUser(username);
            removeClient(clientSocket);
            
            EnhancedMessage leaveMsg(MSG_USER_STATUS, "Server", username + " left the chat");
            broadcastMessage(leaveMsg, -1);
        }
        
        close(clientSocket);
        printLog("Client " + clientIP + " disconnected");
    }
    
    void handleEnhancedMessage(const EnhancedMessage& msg, const std::string& username, int clientSocket) {
        switch (msg.type) {
            case MSG_PRIVATE_CHAT:
                // 私聊消息
                if (!msg.toUser.empty() && !msg.content.empty()) {
                    chatSystem->sendPrivateMessage(username, msg.toUser, msg.content);
                    printLog("Private message from " + username + " to " + msg.toUser + ": " + msg.content);
                }
                break;
                
            case MSG_GROUP_CHAT:
                // 群聊消息
                if (!msg.groupId.empty() && !msg.content.empty()) {
                    chatSystem->sendGroupMessage(msg.groupId, username, msg.content);
                    printLog("Group message in " + msg.groupId + " from " + username + ": " + msg.content);
                }
                break;
                
            case MSG_CREATE_GROUP:
                // 创建群组
                if (!msg.content.empty()) {
                    std::string groupId = chatSystem->createGroup(msg.content, username, "");
                    if (!groupId.empty()) {
                        EnhancedMessage response(MSG_CREATE_GROUP, "Server", "Group created: " + groupId);
                        response.groupId = groupId;
                        sendToClient(clientSocket, response);
                    } else {
                        EnhancedMessage error(MSG_ERROR, "Server", "Failed to create group");
                        sendToClient(clientSocket, error);
                    }
                }
                break;
                
            case MSG_JOIN_GROUP:
                // 加入群组
                if (!msg.groupId.empty()) {
                    if (chatSystem->joinGroup(msg.groupId, username)) {
                        EnhancedMessage response(MSG_JOIN_GROUP, "Server", "Joined group successfully");
                        response.groupId = msg.groupId;
                        sendToClient(clientSocket, response);
                    } else {
                        EnhancedMessage error(MSG_ERROR, "Server", "Failed to join group");
                        sendToClient(clientSocket, error);
                    }
                }
                break;
                
            case MSG_LEAVE_GROUP:
                // 离开群组
                if (!msg.groupId.empty()) {
                    if (chatSystem->leaveGroup(msg.groupId, username)) {
                        EnhancedMessage response(MSG_LEAVE_GROUP, "Server", "Left group successfully");
                        response.groupId = msg.groupId;
                        sendToClient(clientSocket, response);
                    } else {
                        EnhancedMessage error(MSG_ERROR, "Server", "Failed to leave group");
                        sendToClient(clientSocket, error);
                    }
                }
                break;
                
            case MSG_CHAT_HISTORY:
                // 获取聊天记录
                handleChatHistoryRequest(msg, username, clientSocket);
                break;
                
            case MSG_USER_LIST:
                sendUserList(clientSocket);
                break;
                
            case MSG_GROUP_LIST:
                sendGroupList(username, clientSocket);
                break;
                
            case MSG_LOGOUT:
                printLog("User '" + username + "' logged out");
                return; // 退出消息处理循环
                
            default:
                printLog("Unknown message type from " + username + ": " + std::to_string(static_cast<int>(msg.type)));
                break;
        }
        
        // 更新用户活跃时间
        chatSystem->updateUserActivity(username);
    }
    
    void handleChatHistoryRequest(const EnhancedMessage& msg, const std::string& username, int clientSocket) {
        std::vector<EnhancedMessage> messages;
        
        if (msg.chatType == CHAT_PRIVATE && !msg.toUser.empty()) {
            // 获取私聊历史
            messages = chatSystem->getPrivateMessages(username, msg.toUser, 50, 0);
        } else if (msg.chatType == CHAT_GROUP && !msg.groupId.empty()) {
            // 获取群聊历史
            messages = chatSystem->getGroupMessages(msg.groupId, 50, 0);
        }
        
        // 发送历史消息
        for (const auto& historyMsg : messages) {
            sendToClient(clientSocket, historyMsg);
        }
        
        // 发送历史记录结束标志
        EnhancedMessage endMsg(MSG_CHAT_HISTORY, "Server", "History end");
        sendToClient(clientSocket, endMsg);
    }
    
    bool authenticateUser(const std::string& username, int clientSocket, const std::string& clientIP) {
        if (username.empty() || username == "Server") {
            return false;
        }
        
        // 使用数据库持久化系统添加用户
        return chatSystem->addUser(username, clientSocket, clientIP);
    }
    
    void addClient(int socket, const std::string& username, const std::string& ip) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        clients.push_back(std::make_shared<ClientInfo>(socket, username, ip));
    }
    
    void removeClient(int socket) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        clients.erase(
            std::remove_if(clients.begin(), clients.end(),
                [socket](const std::shared_ptr<ClientInfo>& client) {
                    return client->socket == socket;
                }),
            clients.end()
        );
    }
    
    void sendToClient(int socket, const EnhancedMessage& msg) {
        std::string data = msg.serialize();
        send(socket, data.c_str(), data.length(), 0);
    }
    
    void broadcastMessage(const EnhancedMessage& msg, int excludeSocket = -1) {
        std::string data = msg.serialize();
        
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (const auto& client : clients) {
            if (client->socket != excludeSocket) {
                send(client->socket, data.c_str(), data.length(), 0);
            }
        }
    }
    
    void sendUserList(int socket) {
        auto onlineUsers = chatSystem->getOnlineUsers();
        
        std::string userList = "Online users: ";
        for (const auto& user : onlineUsers) {
            userList += user->userId + ", ";
        }
        
        if (userList.length() > 14) { // "Online users: ".length()
            userList = userList.substr(0, userList.length() - 2); // 移除最后的", "
        }
        
        EnhancedMessage listMsg(MSG_USER_LIST, "Server", userList);
        sendToClient(socket, listMsg);
    }
    
    void sendGroupList(const std::string& username, int socket) {
        auto userGroups = chatSystem->getUserGroups(username);
        
        std::string groupList = "Your groups: ";
        for (const auto& group : userGroups) {
            groupList += group->groupName + " (" + group->groupId + "), ";
        }
        
        if (groupList.length() > 13) { // "Your groups: ".length()
            groupList = groupList.substr(0, groupList.length() - 2); // 移除最后的", "
        }
        
        EnhancedMessage listMsg(MSG_GROUP_LIST, "Server", groupList);
        sendToClient(socket, listMsg);
    }
};

// 全局服务器实例用于信号处理
EnhancedChatServer* g_server = nullptr;

void signalHandler(int signal) {
    if (g_server) {
        std::cout << "\nReceived signal " << signal << ", stopping server..." << std::endl;
        g_server->stop();
    }
    exit(0);
}

int main(int argc, char* argv[]) {
    int port = DEFAULT_PORT;
    
    // 解析命令行参数
    if (argc > 1) {
        port = std::atoi(argv[1]);
        if (port <= 0 || port > 65535) {
            std::cerr << "Invalid port number: " << argv[1] << std::endl;
            return 1;
        }
    }
    
    // 设置信号处理
    signal(SIGINT, signalHandler);
    signal(SIGTERM, signalHandler);
    
    // 创建并启动服务器
    EnhancedChatServer server(port);
    g_server = &server;
    
    if (!server.start()) {
        std::cerr << "Failed to start enhanced chat server" << std::endl;
        return 1;
    }
    
    std::cout << "\n=== Enhanced Chat Server with Database ===" << std::endl;
    std::cout << "Features:" << std::endl;
    std::cout << "- Private messaging with history" << std::endl;
    std::cout << "- Group chat with persistent storage" << std::endl;
    std::cout << "- User management with MySQL database" << std::endl;
    std::cout << "- Message history and user activity tracking" << std::endl;
    std::cout << "============================================\n" << std::endl;
    
    // 运行服务器
    server.run();
    
    return 0;
} 